#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/CameraInfo.h>
#include <geometry_msgs/PoseStamped.h>
#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/calib3d.hpp>
#include <opencv2/aruco.hpp> // For ArUco and ChArUco
#include <opencv2/highgui.hpp> // For imshow, waitKey (optional debug)
#include <opencv2/aruco/charuco.hpp> // For ChArUco board detection

// For tf2 conversions
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h> // For toMsg

#include <iostream>
#include <opencv2/core/version.hpp>


class CharucoDetector {
public:
    CharucoDetector() : nh_(), pnh_("~"), it_(nh_), cam_info_received_(false) {
        // Get ChArUco board parameters
        pnh_.param("squares_x", squares_x_, 7); // Number of chessboard squares in X direction
        pnh_.param("squares_y", squares_y_, 5); // Number of chessboard squares in Y direction
        pnh_.param("square_length_m", square_length_m_, 0.04); // Side length of a chessboard square in meters
        pnh_.param("marker_length_m", marker_length_m_, 0.02); // Side length of the ArUco marker in meters
        pnh_.param("dictionary_id", dictionary_id_, (int)cv::aruco::DICT_6X6_250); // ArUco dictionary ID

        // Get common parameters
        pnh_.param("image_is_rectified", image_is_rectified_, true);
        std::string image_topic, camera_info_topic, pose_topic_name;
        pnh_.param<std::string>("image_topic", image_topic, "/camera/color/image_raw"); // Default to raw, adjust as needed
        pnh_.param<std::string>("camera_info_topic", camera_info_topic, "/camera/color/camera_info");
        pnh_.param<std::string>("pose_topic", pose_topic_name, "/charuco_pose");
        pnh_.param<bool>("show_debug_image", show_debug_image_, true);
        pnh_.param("axis_length_m", axis_length_m_, 0.05); // Length of the drawn axis for visualization

        // Create ArUco dictionary
        try {
            dictionary_ = cv::aruco::getPredefinedDictionary(static_cast<cv::aruco::PREDEFINED_DICTIONARY_NAME>(dictionary_id_));
        } catch (const cv::Exception& e) {
            ROS_FATAL("Failed to get predefined dictionary with ID %d. Error: %s", dictionary_id_, e.what());
            ROS_FATAL("Valid IDs are e.g., DICT_6X6_250=10, DICT_5X5_250=8, DICT_4X4_50=0 etc.");
            ros::shutdown();
            return;
        }


        // Create ChArUco board
        charuco_board_ = cv::aruco::CharucoBoard::create(squares_x_, squares_y_,
                                                         (float)square_length_m_,
                                                         (float)marker_length_m_,
                                                         dictionary_);

        // Create ArUco detector parameters
        detector_params_ = cv::aruco::DetectorParameters::create();
        // You can refine parameters here, e.g.:
        // detector_params_->cornerRefinementMethod = cv::aruco::CORNER_REFINE_SUBPIX;
        // detector_params_->adaptiveThreshWinSizeMin = 3;
        // detector_params_->adaptiveThreshWinSizeMax = 23;
        // detector_params_->adaptiveThreshConstant = 7;


        // Subscribers
        image_sub_ = it_.subscribe(image_topic, 1, &CharucoDetector::imageCallback, this);
        cam_info_sub_ = nh_.subscribe(camera_info_topic, 1, &CharucoDetector::cameraInfoCallback, this);

        // Publisher
        pose_pub_ = nh_.advertise<geometry_msgs::PoseStamped>(pose_topic_name, 10);
        if (show_debug_image_) {
            debug_image_pub_ = it_.advertise("/charuco_detector/debug_image", 1);
            cv::namedWindow("ChArUco Detection", cv::WINDOW_AUTOSIZE);
        }

        ROS_INFO("ChArUco detector initialized.");
        ROS_INFO("Board: %d x %d squares", squares_x_, squares_y_);
        ROS_INFO("Square length: %.3f m, Marker length: %.3f m", square_length_m_, marker_length_m_);
        ROS_INFO("Dictionary ID: %d", dictionary_id_);
        ROS_INFO("Subscribing to image: %s", image_sub_.getTopic().c_str());
        ROS_INFO("Subscribing to camera info: %s", cam_info_sub_.getTopic().c_str());
        ROS_INFO("Publishing pose to: %s", pose_pub_.getTopic().c_str());
        ROS_INFO("Image is rectified: %s", image_is_rectified_ ? "true" : "false");
    }

    ~CharucoDetector() {
        if (show_debug_image_) {
            cv::destroyWindow("ChArUco Detection");
        }
    }

private:
    void cameraInfoCallback(const sensor_msgs::CameraInfoConstPtr& msg) {
        if (cam_info_received_ && image_is_rectified_) { // If rectified, cam_info should be static
             // Optionally shut down if info is static and you only need it once.
             // cam_info_sub_.shutdown();
             // ROS_INFO("Camera info subscriber shut down (rectified image).");
            return;
        }
         if (cam_info_received_ && !image_is_rectified_ && msg->D.empty()){
            // This case might happen if camera_info for a raw stream incorrectly reports no distortion.
            // Or if you expect distortion but it's not provided.
            ROS_WARN_THROTTLE(5.0, "Received camera_info for non-rectified image, but distortion coefficients (D) are empty. Assuming zero distortion.");
        }


        camera_matrix_ = cv::Mat(3, 3, CV_64F, (void*)msg->K.data()).clone();
        if (image_is_rectified_) {
            dist_coeffs_ = cv::Mat::zeros(1, 5, CV_64F); // Assume 5 coeffs for simplicity, can be 1xN
        } else {
            if (!msg->D.empty()) {
                dist_coeffs_ = cv::Mat(msg->D.size(), 1, CV_64F, (void*)msg->D.data()).clone();
            } else {
                // If D is empty for a non-rectified image, assume no distortion.
                // This might be an issue with the camera_info publisher.
                dist_coeffs_ = cv::Mat::zeros(1, 5, CV_64F);
            }
        }
        cam_info_received_ = true;
        ROS_INFO("Camera info received. Distortion coeffs size: %ld", msg->D.size());
    }

    void imageCallback(const sensor_msgs::ImageConstPtr& msg) {
        if (!cam_info_received_) {
            ROS_WARN_THROTTLE(5.0, "Waiting for camera info...");
            return;
        }

        cv_bridge::CvImagePtr cv_ptr;
        try {
            cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
            return;
        }

        cv::Mat image_display = cv_ptr->image.clone(); // For drawing visualizations

        std::vector<int> marker_ids;
        std::vector<std::vector<cv::Point2f>> marker_corners;

        // Detect ArUco markers
        cv::aruco::detectMarkers(cv_ptr->image, dictionary_, marker_corners, marker_ids, detector_params_);

        if (!marker_ids.empty()) {
            if (show_debug_image_ || debug_image_pub_.getNumSubscribers() > 0) {
                cv::aruco::drawDetectedMarkers(image_display, marker_corners, marker_ids);
            }

            std::vector<cv::Point2f> charuco_corners;
            std::vector<int> charuco_ids;
            // Interpolate ChArUco corners
            // Provide cameraMatrix and distCoeffs for better accuracy if image is not rectified
            cv::aruco::interpolateCornersCharuco(marker_corners, marker_ids, cv_ptr->image, charuco_board_,
                                                 charuco_corners, charuco_ids, camera_matrix_, dist_coeffs_);

            if (!charuco_ids.empty()) {
                 if (show_debug_image_ || debug_image_pub_.getNumSubscribers() > 0) {
                    cv::aruco::drawDetectedCornersCharuco(image_display, charuco_corners, charuco_ids, cv::Scalar(255, 0, 0));
                }

                cv::Vec3d rvec, tvec; // cv::Vec3d for rvec, tvec with aruco functions
                bool pose_valid = cv::aruco::estimatePoseCharucoBoard(charuco_corners, charuco_ids, charuco_board_,
                                                                    camera_matrix_, dist_coeffs_, rvec, tvec);

                if (pose_valid) {
                    if (show_debug_image_ || debug_image_pub_.getNumSubscribers() > 0) {
                        cv::drawFrameAxes(image_display, camera_matrix_, dist_coeffs_, rvec, tvec, (float)axis_length_m_);
                    }

                    geometry_msgs::PoseStamped pose_stamped_msg;
                    pose_stamped_msg.header.stamp = msg->header.stamp;
                    pose_stamped_msg.header.frame_id = msg->header.frame_id; // Should be camera_optical_frame

                    pose_stamped_msg.pose.position.x = tvec[0];
                    pose_stamped_msg.pose.position.y = tvec[1];
                    pose_stamped_msg.pose.position.z = tvec[2];

                    cv::Mat R_matrix;
                    cv::Mat rvec_mat = (cv::Mat_<double>(3,1) << rvec[0], rvec[1], rvec[2]); // Convert Vec3d to Mat
                    cv::Rodrigues(rvec_mat, R_matrix);

                    tf2::Matrix3x3 tf_R_matrix(
                        R_matrix.at<double>(0,0), R_matrix.at<double>(0,1), R_matrix.at<double>(0,2),
                        R_matrix.at<double>(1,0), R_matrix.at<double>(1,1), R_matrix.at<double>(1,2),
                        R_matrix.at<double>(2,0), R_matrix.at<double>(2,1), R_matrix.at<double>(2,2)
                    );
                    tf2::Quaternion tf_q;
                    tf_R_matrix.getRotation(tf_q);
                    tf_q.normalize();
                    pose_stamped_msg.pose.orientation = tf2::toMsg(tf_q);

                    pose_pub_.publish(pose_stamped_msg);
                } else {
                    ROS_WARN_THROTTLE(1.0, "ChArUco pose estimation failed.");
                }
            } else {
                 ROS_WARN_THROTTLE(1.0, "Not enough ChArUco corners found for pose estimation (IDs count: %zu).", charuco_ids.size());
            }
        } else {
            ROS_WARN_THROTTLE(5.0, "No ArUco markers detected in the image.");
        }

        if (show_debug_image_) {
            cv::imshow("ChArUco Detection", image_display);
            cv::waitKey(1);
        }
        if (debug_image_pub_.getNumSubscribers() > 0) {
            cv_bridge::CvImage out_msg;
            out_msg.header = msg->header;
            out_msg.encoding = sensor_msgs::image_encodings::BGR8;
            out_msg.image = image_display;
            debug_image_pub_.publish(out_msg.toImageMsg());
        }
    }

    ros::NodeHandle nh_;
    ros::NodeHandle pnh_;
    image_transport::ImageTransport it_;

    image_transport::Subscriber image_sub_;
    ros::Subscriber cam_info_sub_;
    ros::Publisher pose_pub_;
    image_transport::Publisher debug_image_pub_;

    // ChArUco parameters
    int squares_x_;
    int squares_y_;
    double square_length_m_;
    double marker_length_m_;
    int dictionary_id_;
    double axis_length_m_;

    cv::Ptr<cv::aruco::CharucoBoard> charuco_board_;
    cv::Ptr<cv::aruco::Dictionary> dictionary_;
    cv::Ptr<cv::aruco::DetectorParameters> detector_params_;

    // Camera parameters
    cv::Mat camera_matrix_;
    cv::Mat dist_coeffs_;
    bool cam_info_received_;
    bool image_is_rectified_;

    bool show_debug_image_;
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "charuco_detector_node");
    CharucoDetector detector;
    ros::spin();
    return 0;
}